package com.tencent.welink.util;

import com.tencent.welink.sm.*;
import org.apache.commons.lang3.StringUtils;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import sun.misc.BASE64Decoder;
import sun.misc.BASE64Encoder;

import java.util.regex.Matcher;
import java.util.regex.Pattern;

/**
 * Created by junchihe on 2018/12/7.
 */
public class SM4_CBC implements IEncryptAndDecrypt {
    private static final Logger log = LogManager.getLogger(SM4_CBC.class);



    private final String iv ;
    private  SM4  sm4_encrypt ;
    private  SM4_Context ctx_encrypt;
    private  SM4  sm4_decrypt ;
    private  SM4_Context ctx_decrypt;

    public SM4_CBC(String secretKey,String iv,boolean en,boolean de) {
        this.iv = iv;
        try {
            if (en) {
                ctx_encrypt = new SM4_Context();
                ctx_encrypt.isPadding = true;
                ctx_encrypt.mode = SM4.SM4_ENCRYPT;
                sm4_encrypt = new SM4();
                sm4_encrypt.sm4_setkey_enc(ctx_encrypt, secretKey.getBytes());
            }

            if (de) {
                ctx_decrypt = new SM4_Context();
                ctx_decrypt.isPadding = true;
                ctx_decrypt.mode = SM4.SM4_DECRYPT;
                sm4_decrypt = new SM4();
                sm4_decrypt.sm4_setkey_dec(ctx_decrypt, secretKey.getBytes());
            }
        } catch (Exception e) {
            log.error("SM4 inital exception",e);
        }
    }

    @Override
    public byte[] encrypt(byte[] content) {
        try {
            byte[] encrypted = sm4_encrypt.sm4_crypt_cbc(ctx_encrypt, iv.getBytes(), content);
            ;
            if(encrypted != null && encrypted.length>0)
                return encrypted;
        } catch (Exception e) {
            log.error(e);
        }
        return null;
    }

    @Override
    public String encrypt(String content) {
        try {
            if(StringUtils.isBlank(content)) return content;

            byte[] encrypted = encrypt(content.getBytes("utf-8"));
            if(encrypted != null && encrypted.length>0) {
                String cipherText = new BASE64Encoder().encode(encrypted);
                if (cipherText != null && cipherText.trim().length() > 0) {
                    Pattern p = Pattern.compile("\\s*|\t|\r|\n");
                    Matcher m = p.matcher(cipherText);
                    cipherText = m.replaceAll("");
                    return cipherText;
                }
            }
        } catch (Exception e) {
            log.error(e);
        }
        return null;
    }

    @Override
    public String decrypt(String content) {
        try {
            if(StringUtils.isBlank(content)) return content;

            byte[] decrypted = decrypt(new BASE64Decoder().decodeBuffer(content));
            if(decrypted != null &&  decrypted.length>0)
                return new String(decrypted, "utf-8");
        } catch (Exception e) {
            log.error(e);
        }
        return null;
    }

    @Override
    public byte[] decrypt(byte[] content) {
        try {
            if(content == null || content.length<=0) return content;

            byte[] decrypted = sm4_decrypt.sm4_crypt_cbc(ctx_decrypt, iv.getBytes(), content);
            if(decrypted!= null && decrypted.length>0)
                return decrypted;
        } catch (Exception e) {
            log.error(e);

        }
        return content;
    }


    public static void main(String[] args) {
        SM4_CBC sm4_cbc = new SM4_CBC("1234567812345678","1234567812345678",true,true);
        String de = "12121212545455566";
        System.out.println(de);
        String en = sm4_cbc.encrypt(de);
        System.out.println(en);
        System.out.println(sm4_cbc.decrypt(en));
    }
}
